package developer.zyk.ability.chat.ability.rag.vector;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import developer.zyk.ability.knowledge.dao.entity.ChatKnowledgeRecordDetailDO;
import developer.zyk.ability.knowledge.dao.mappers.ChatKnowledgeRecordDetailMapper;
import jakarta.annotation.Resource;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.context.event.ApplicationStartedEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Objects;

/**
 * @author zhangyongkang
 * @time 2025/5/16 17:16
 * @description
 */
@Component
public class VectorStoreManager implements ApplicationListener<ApplicationStartedEvent> {
    @Resource
    private UserStoreVector vectorStore;
    @Resource
    private ChatKnowledgeRecordDetailMapper chatKnowledgeRecordDetailMapper;

    public void loadAll() {
        vectorStore.loadAllOnStart();
    }

    public void load(List<Long> recordIds) {
        if (CollUtil.isEmpty(recordIds)) {
            return;
        }
        vectorStore.load(recordIds);
    }


    public void removeByRecordIds(
            List<Long> recordIds
    ) {
        if (CollUtil.isEmpty(recordIds)) {
            return;
        }
        List<String> documentIds = chatKnowledgeRecordDetailMapper.selectList(
                        new QueryWrapper<ChatKnowledgeRecordDetailDO>()
                                .lambda()
                                .in(ChatKnowledgeRecordDetailDO::getRecordId, recordIds)
                ).stream()
                .filter(Objects::nonNull)
                .map(ChatKnowledgeRecordDetailDO::getDocumentOutId)
                .filter(Objects::nonNull)
                .toList();

        if (CollUtil.isEmpty(documentIds)) {
            return;
        }
        ((VectorStore) vectorStore).delete(documentIds);
    }


    public void realDelete(
            List<Long> recordIds
    ) {
        if (CollUtil.isEmpty(recordIds)) {
            return;
        }
        vectorStore.realDelete(recordIds);
    }

    @Override
    public void onApplicationEvent(ApplicationStartedEvent event) {
        loadAll();
    }


}
